import os

from zhipuai import ZhipuAI


# 初始化API客户端
def initialize_client(api_key):
    """
    初始化智谱AI API客户端。
    参数:
    - api_key (str): 智谱AI提供的API密钥。
    返回:
    - client: 初始化后的智谱AI API客户端对象。
    """
    return ZhipuAI(api_key=api_key)


# 多轮对话函数
def multi_round_chat(api_key):
    """
    开启基于智谱AI的多轮对话。
    本函数使用默认参数与用户进行多轮对话，使用glm-4模型。
    用户输入“退出”时，结束对话。
    参数:
    - api_key (str): 智谱AI提供的API密钥。
    返回:
    - None
    """
    client = initialize_client(api_key)
    messages = []  # 初始化消息队列
    while True:
        user_input = input("用户: ")
        if user_input.lower() == "退出":  # 检查用户是否输入退出指令
            print("对话已结束。")
            break
        messages.append({"role": "user", "content": user_input})
        response = client.chat.completions.create(
            model="glm-4",
            messages=messages,
            # 其他参数使用默认值
        )
        assistant_message = response.choices[0].message
        messages.append({"role":"assistant","content":assistant_message.content})  # 将助手的回复添加到消息队列中
        print(f"助手: {assistant_message.content}")


# 函数说明文档
multi_round_chat.__doc__ += """
函数使用示例:
    # 假设已经有一个有效的API密钥字符串赋值给了变量 `api_key`
    multi_round_chat(api_key)

    # 在对话中，用户可以输入问题或者指令
    用户: 你好！
    助手: 你好！有什么可以帮助你的吗？

    # 当用户想要退出对话时，输入“退出”
    用户: 退出
    助手: 对话已结束。

注意：
- 本函数需要先安装智谱AI提供的Python SDK。
- 函数内部已默认设置了模型为glm-4，其他参数使用默认值。
- 本函数无返回值，对话内容通过标准输出打印。
"""

if __name__ == "__main__":
    api_key = os.environ["ZHIPU_API_KEY"]
    multi_round_chat(api_key)